// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/vector/vector_maker.h>

namespace kumo::pollux::test {

class MayHaveNullsRecursiveTest : public testing::Test {
 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }

  std::shared_ptr<pollux::memory::MemoryPool> pool_{
      memory::memoryManager()->addLeafPool()};
  VectorMaker vectorMaker_{pool_.get()};

  enum class TestOptions {
    WITH_NULL_CONTAINERS,
    WITH_NULL_VALUES,
  };

  std::vector<std::optional<int32_t>> generateIntInput() {
    return generateIntInput(false);
  }

  std::vector<std::optional<int32_t>> generateIntInputWithNulls() {
    return generateIntInput(true);
  }

  ArrayVectorPtr generateArrayVector(
      const std::unordered_set<TestOptions>& testOptions = {}) {
    std::function<bool(vector_size_t)> noNulls = nullptr;

    return vectorMaker_.array_vector<int32_t>(
        10,
        [](vector_size_t row) { return row; },
        [](vector_size_t idx) { return idx; },
        testOptions.count(TestOptions::WITH_NULL_CONTAINERS) > 0
            ? [](vector_size_t row) { return row % 5 == 0; }
            : noNulls,
        testOptions.count(TestOptions::WITH_NULL_VALUES) > 0
            ? [](vector_size_t idx) { return idx % 7 == 0; }
            : noNulls);
  }

  MapVectorPtr generateMapVector(
      const std::unordered_set<TestOptions>& testOptions = {}) {
    std::function<bool(vector_size_t)> noNulls = nullptr;

    return vectorMaker_.map_vector<int32_t, int32_t>(
        10,
        [](vector_size_t row) { return row; },
        [](vector_size_t idx) { return idx; },
        [](vector_size_t idx) { return idx % 10; },
        testOptions.count(TestOptions::WITH_NULL_CONTAINERS) > 0
            ? [&](vector_size_t row) { return row % 5 == 0; }
            : noNulls,
        testOptions.count(TestOptions::WITH_NULL_VALUES) > 0
            ? [&](vector_size_t row) { return row % 7 == 0; }
            : noNulls);
  }

 private:
  std::vector<std::optional<int32_t>> generateIntInput(bool withNulls) {
    std::vector<std::optional<int32_t>> input;
    for (int32_t i = 0; i < 10; ++i) {
      if (withNulls && i % 5 == 0) {
        input.push_back(std::nullopt);
      } else {
        input.push_back(i);
      }
    }

    return input;
  }
};

TEST_F(MayHaveNullsRecursiveTest, biasHasNulls) {
  auto bias_vector = vectorMaker_.bias_vector(generateIntInputWithNulls());

  ASSERT_TRUE(bias_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, biasNullFree) {
  auto bias_vector = vectorMaker_.bias_vector(generateIntInput());

  ASSERT_FALSE(bias_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, arrayHasNulls) {
  auto array_vector = generateArrayVector({TestOptions::WITH_NULL_CONTAINERS});

  ASSERT_TRUE(array_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, arrayHasNullElements) {
  auto array_vector = generateArrayVector({TestOptions::WITH_NULL_VALUES});

  ASSERT_TRUE(array_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, arrayNullFree) {
  auto array_vector = generateArrayVector();

  ASSERT_FALSE(array_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, mapHasNulls) {
  auto map_vector = generateMapVector({TestOptions::WITH_NULL_CONTAINERS});

  ASSERT_TRUE(map_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, mapHasNullValues) {
  auto map_vector = generateMapVector({TestOptions::WITH_NULL_VALUES});

  ASSERT_TRUE(map_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, mapNullFree) {
  auto map_vector = generateMapVector();

  ASSERT_FALSE(map_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, rowHasNulls) {
  auto row_vector = vectorMaker_.row_vector(
      {vectorMaker_.flat_vector_nullable(generateIntInput()),
       generateArrayVector(),
       generateMapVector()});

  row_vector->set_null(3, true);

  ASSERT_TRUE(row_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, rowChildrenHaveNulls) {
  auto row_vector = vectorMaker_.row_vector(
      {vectorMaker_.flat_vector_nullable(generateIntInputWithNulls()),
       generateArrayVector(),
       generateMapVector()});

  ASSERT_TRUE(row_vector->may_have_nulls_recursive());

  row_vector = vectorMaker_.row_vector(
      {vectorMaker_.flat_vector_nullable(generateIntInput()),
       generateArrayVector({TestOptions::WITH_NULL_VALUES}),
       generateMapVector()});

  ASSERT_TRUE(row_vector->may_have_nulls_recursive());

  row_vector = vectorMaker_.row_vector(
      {vectorMaker_.flat_vector_nullable(generateIntInput()),
       generateArrayVector(),
       generateMapVector({TestOptions::WITH_NULL_CONTAINERS})});

  ASSERT_TRUE(row_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, rowNullFree) {
  auto row_vector = vectorMaker_.row_vector(
      {vectorMaker_.flat_vector_nullable(generateIntInput()),
       generateArrayVector(),
       generateMapVector()});

  ASSERT_FALSE(row_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, constantNull) {
  auto constant_vector = vectorMaker_.constant_vector<int32_t>({std::nullopt});

  ASSERT_TRUE(constant_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, constantNullFree) {
  auto constant_vector = vectorMaker_.constant_vector<int32_t>({1});

  ASSERT_FALSE(constant_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, dictionaryHasNulls) {
  auto dictionary_vector =
      vectorMaker_.dictionary_vector(generateIntInputWithNulls());

  ASSERT_TRUE(dictionary_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, dictionaryNullFree) {
  auto dictionary_vector = vectorMaker_.dictionary_vector(generateIntInput());

  ASSERT_FALSE(dictionary_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, flatHasNulls) {
  auto flat_vector =
      vectorMaker_.flat_vector_nullable(generateIntInputWithNulls());

  ASSERT_TRUE(flat_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, flatNullFree) {
  auto flat_vector = vectorMaker_.flat_vector_nullable(generateIntInput());

  ASSERT_FALSE(flat_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, lazyHasNulls) {
  auto lazyVector = vectorMaker_.lazy_flat_vector<int32_t>(
      10,
      [](vector_size_t row) { return row; },
      [](vector_size_t row) { return row % 3 == 0; });

  ASSERT_TRUE(lazyVector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, lazyNullFree) {
  auto lazyVector = vectorMaker_.lazy_flat_vector<int32_t>(
      10, [](vector_size_t row) { return row; });

  ASSERT_FALSE(lazyVector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, sequenceHasNulls) {
  auto sequence_vector =
      vectorMaker_.sequence_vector(generateIntInputWithNulls());

  ASSERT_TRUE(sequence_vector->may_have_nulls_recursive());
}

TEST_F(MayHaveNullsRecursiveTest, sequenceNullFree) {
  auto sequence_vector = vectorMaker_.sequence_vector(generateIntInput());

  ASSERT_FALSE(sequence_vector->may_have_nulls_recursive());
}
} // namespace kumo::pollux::test
